from langchain_core.runnables import RunnableLambda
import asyncio

# 定义一个函数
def add_one(x: int) -> int:
    return x + 1


runnable = RunnableLambda(add_one)

# 同步调用
print(runnable.invoke(1))
print(runnable.batch([1, 2, 3, 4, 5]))


# 异步调用
async def add_one_async(x: int) -> int:
    return x + 1


async def main():
    # 直接异步调用
    print(await runnable.ainvoke(1))
    print(await runnable.abatch([1, 2, 3, 4, 5]))

    # 添加异步函数 再异步调用
    runnable_async = RunnableLambda(add_one, afunc=add_one_async)
    print(runnable_async.invoke(1))
    print(runnable_async.batch([1, 2, 3, 4, 5]))

# 执行异步函数
asyncio.run(main())